import cv2
import os
import functools
import tkinter as tk
import tkinter.messagebox

'''
图像处理类, 主要负责图像的读取和ROI的选择
后续会增加数据增强的接口
'''


class ImageReader():
    def __init__(self, videoPath=None, imagePath=None, trackerType='KCF'):
        '''
        :param videoPath: 视频路径
        :param imagePath: 图像路径
        :param trackerType: 追踪器类型
        '''
        # 标识图像读取状态
        self.status = True
        if videoPath:
            try:
                self.cap = cv2.VideoCapture(videoPath)
                self.getImage = self.fromVideo
            except:
                self.status = False
            else:
                self.status = True
        if imagePath:
            try:
                self.imagesPath = imagePath
                self.getImage = self.fromImage
            except:
                self.status = False
            else:
                self.status = True
        self.trackerType = trackerType
        # 记录类型名与追踪器对应关系的字典
        self.trackers = {'BOOSTING': cv2.TrackerBoosting_create, 'MIL': cv2.TrackerMIL_create,
                         'KCF': cv2.TrackerKCF_create, 'TLD': cv2.TrackerTLD_create,
                         'MEDIANFLOW': cv2.TrackerMedianFlow_create, 'CSRT': cv2.TrackerCSRT_create,
                         'MOSSE': cv2.TrackerMOSSE_create}
        # 追踪器
        self.tracker = None
        # 类名
        self.className = None
        # 获取类名的输入框
        self.classNameEntry = None
        self.window = None
        self.font = ('Arial', 15)

    # 确认键 保存类型名
    def confirm(self):
        self.className = self.classNameEntry.get()
        self.window.destroy()

    # 从视频中获取图像, 迭代器方式
    def fromVideo(self):
        while True:
            _, frame = self.cap.read()
            yield frame

    # 从图像文件夹中获取图像, 迭代器方式
    def fromImage(self):
        imageNames = [i for i in os.listdir(self.imagesPath) if
                      i.split('.')[-1].lower() in ['jpg', 'jpeg', 'png', 'bmp', 'gif']]
        for im in imageNames:
            yield cv2.imread(os.path.join(self.imagesPath, im))

    # 统一接口, 对外都通过此方法获取图像
    def getImage(self):
        pass

    # 获取ROI
    def readBbox(self, frame):
        '''
        :param frame: 需要获取ROI的图像
        :return: ROI和类别名
        '''
        lastName = self.className
        self.className = None
        bbox = None
        # 防止用户输入的类别为空
        while not self.className:
            bbox = cv2.selectROI(frame, False)
            if bbox == (0, 0, 0, 0):
                break
            # 创建窗口要求用户输入类型名
            self.window = tk.Tk()
            tk.Label(self.window, text='请输入类别名', width=12, font=self.font).grid(row=0)
            self.classNameEntry = tk.Entry(self.window, width=12, font=self.font)
            if lastName:
                self.classNameEntry.insert('end', lastName)
            self.classNameEntry.grid(row=1)
            tk.Button(self.window, text='确认', command=self.confirm, font=self.font).grid(row=2)
            self.window.mainloop()
        return bbox, self.className

    # 单目标追踪器, 已弃用, 统一使用多目标追踪器
    def initTrackerSingle(self, frame):
        bbox, className = self.readBbox(frame)
        self.tracker = self.trackers[self.trackerType]()
        self.tracker.init(frame, bbox)
        return bbox, className

    # 获取多个ROI
    def getBboxs(self, frame):
        '''
        :param frame: 需要获取ROI的图像
        :return: ROI和其名字
        '''
        self.more = True
        bboxs = []
        names = []

        while True:
            bbox, name = self.readBbox(frame)
            if bbox[2] * bbox[3] >= 25:  # 忽略太小的ROI
                bboxs.append(bbox)
                names.append(name)
            if not tk.messagebox.askyesno(title='标注完成', message='是否继续标注?'):
                # print('finish label')
                break
        return bboxs, names

    # 初始化跟踪器
    def initTrackerMulti(self, frame):
        '''
        :param frame: 需要获取ROI的图像
        :return: ROI和名字
        '''
        bboxs, names = self.getBboxs(frame)
        self.tracker = cv2.MultiTracker_create()
        for bbox in bboxs:
            self.tracker.add(self.trackers[self.trackerType](), frame, bbox)
        # print('init tracker successfully!')
        return bboxs, names


if __name__ == '__main__':
    ir = ImageReader(videoPath='/home/ningfeng/视频/小陀螺1.avi')
    frameReader = ir.getImage()
    while True:
        frame = next(frameReader)
        cv2.imshow('name', frame)
        cv2.waitKey(10)
